/* TCP Implemenation */

/* Copyright (C) 2012 David Zanetti
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; version 2 of the License.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License along
 *  with this program; if not, write to the Free Software Foundation, Inc.,
 *  51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 */

#include <avr/io.h>
#include <avr/pgmspace.h>
#include <stdlib.h>
#include <string.h>
#include "global.h"
#include "tcp.h"
#include "encx24j600.h"
#include "serial.h"

/* TCP header format */

typedef struct {
		uint16_t sport;
		uint16_t dport;
		uint32_t seq;
		uint32_t ack;
		uint16_t control;
		uint16_t window;
		uint16_t crc;
		uint16_t urg;
} tcp_header_t;

/* TCP header flags */
#define TCP_URG (1 << 5)
#define TCP_ACK (1 << 4)
#define TCP_PSH (1 << 3)
#define TCP_RST (1 << 2)
#define TCP_SYN (1 << 1)
#define TCP_FIN (1 << 0)

/* TCP connection table */
typedef struct {
		uint8_t src[4];
		uint16_t sport;
		uint16_t dport;
		uint8_t state;
		uint32_t last_ack;
		uint32_t last_seq;
		uint32_t activity; /* when we last got a packet */
} tcp_conntable_t;

/* TCP states for the connection table
   This does not include TIMEWAIT or SYN-SENT since we are listen only
   "CLOSED" is really unused, since listening is a hardcode state */
#define TCP_STATE_CLOSED 0
#define TCP_STATE_SYN_REC 1
#define TCP_STATE_ESTABLISHED 2
#define TCP_STATE_CLOSE_WAIT 3
#define TCP_STATE_LAST_ACK 4
#define TCP_STATE_FINWAIT1 5
#define TCP_STATE_FINWAIT2 6
#define TCP_STATE_CLOSING 7

/* size of the connection table */
#define TCP_CONNTABLE_SIZE 4

/* the connection table */
tcp_conntable_t tcp_conntable[TCP_CONNTABLE_SIZE];

/* Prototypes */
uint16_t tcp_psuedocrc(uint16_t length, uint8_t *src, uint8_t *dst, uint8_t *pkt);
uint8_t tcp_conntable_new(tcp_header_t pkt, uint8_t *src);
void tcp_conntable_tick(void);
void tcp_makerst(tcp_header_t *pkt, tcp_header_t *new, uint8_t *src);

/* Process incoming TCP packets */
uint8_t tcp_process(uint8_t slot, uint8_t *src, uint8_t *dst) {
	tcp_header_t pkt, new;
	uint16_t dlen, hlen, crc;
	uint8_t tcp_flags;
	uint8_t n, id;

	/* retrieve the TCP header */
	if (encx24j600_rx(&pkt, sizeof(tcp_header_t)) != sizeof(tcp_header_t)) {
		serial0_tx_PGM(PSTR("tcp: short frame, discarding\r\n"));
		return E_NOTHING;
	}

	/* extract flags */
	tcp_flags = (swap16(pkt.control)) & 0x3f;

	/* header length */
	hlen = (swap16(pkt.control) >> 12) * 4;

	/* TCP doesn't have a length field, so compute it from unread size plus the size of the header */
	dlen = encx24j600_rx_unread() - (hlen - sizeof(tcp_header_t));

#ifdef DEBUG_TCP
	serial0_tx_PGM(PSTR("tcp: rx "));
	if (tcp_flags & TCP_URG) {
		serial0_tx_cout('U');
	}
	if (tcp_flags & TCP_ACK) {
		serial0_tx_cout('A');
	}
	if (tcp_flags & TCP_PSH) {
		serial0_tx_cout('P');
	}
	if (tcp_flags & TCP_RST) {
		serial0_tx_cout('R');
	}
	if (tcp_flags & TCP_SYN) {
		serial0_tx_cout('S');
	}
	if (tcp_flags & TCP_FIN) {
		serial0_tx_cout('F');
	}
	serial0_tx_PGM(PSTR(" sport "));
	serial0_tx_dec(swap16(pkt.sport));
	serial0_tx_PGM(PSTR(" dport "));
	serial0_tx_dec(swap16(pkt.dport));
	serial0_tx_PGM(PSTR(" crc "));
	serial0_tx_hex(pkt.crc >> 8);
	serial0_tx_hex(pkt.crc & 0xff);
	serial0_tx_PGM(PSTR(" hlen "));
	serial0_tx_dec(hlen);
	serial0_tx_PGM(PSTR(" dlen "));
	serial0_tx_dec(dlen);
	serial0_tx_cr();
#endif

	/* compute checksum */
	/* psuedoheader and fixed TCP header first */
	crc = tcp_psuedocrc(dlen+hlen, src, dst, (uint8_t *) &pkt);
	/* unread data. This includes option headers, hence the jigging the length */
	crc = encx24j600_rx_crc_unread(dlen + (hlen - sizeof(tcp_header_t)), swap16(crc));

#ifdef DEBUG_TCP
	serial0_tx_PGM(PSTR("tcp: computed crc "));
	serial0_tx_hex(crc >> 8);
	serial0_tx_hex(crc & 0xff);
	serial0_tx_cr();
#endif

	if (crc != 0) {
		serial0_tx_PGM(PSTR("tcp: dropping corrupt packet\r\n"));
		return E_NOTHING;
	}

	/* fixme: handle tcp options */
	/* skip options */
	encx24j600_rx_skip(hlen-sizeof(tcp_header_t));

	/* if this is a packet related to an existing connection, then find it */
	for (n = 0; n < TCP_CONNTABLE_SIZE; n++) {
		/* compare source IP address, sport, and dport to table */
		if (tcp_conntable[n].state != TCP_STATE_CLOSED &&
			src[0] == tcp_conntable[n].src[0] &&
			src[1] == tcp_conntable[n].src[1] &&
			src[2] == tcp_conntable[n].src[2] &&
			src[3] == tcp_conntable[n].src[3] &&
			swap16(pkt.sport) == tcp_conntable[n].sport &&
			swap16(pkt.dport) == tcp_conntable[n].dport) {
#ifdef DEBUG_TCP
			serial0_tx_PGM(PSTR("tcp: found match for existing connection in table#"));
			serial0_tx_dec(n);
			serial0_tx_cr();
#endif

			/* reject it for now, since we're just testing */
			tcp_makerst(&pkt,&new,src);

			/* store our header and pass it up the stack */
			if (!encx24j600_tx_prepend(slot,(uint8_t *)&new,sizeof(tcp_header_t))) {
				serial0_tx_PGM(PSTR("tcp: failed to prepend our header\r\n"));
				return E_NOTHING;
			}

			return E_REPLY; /* we should probably handle it really tho */
		}
	}

	/* this is not a packet for an associated connection. if it's got anything but SYN set, dump it */
	if (tcp_flags != TCP_SYN) {
#ifdef DEBUG_TCP
		serial0_tx_PGM(PSTR("tcp: unknown connection, dropping packet\r\n"));
#endif
	}

	/* this is a SYN packet, let see if we're listening for it */

	switch(swap16(pkt.dport)) {
		case 80:
			/* HTTP */
			id = tcp_conntable_new(pkt,src);
			if (id != 0xff) {
				/* do nothing for now */
				serial0_tx_PGM(PSTR("tcp: accept tcp/80 table#"));
				serial0_tx_dec(id);
				serial0_tx_cr();
			} else {
				serial0_tx_PGM(PSTR("tcp: reject, no table space\r\n"));
			}
			break;
	}

	/* whatever this was, reset it */
	tcp_makerst(&pkt,&new,src);

	/* store our header and pass it up the stack */
	if (!encx24j600_tx_prepend(slot,(uint8_t *)&new,sizeof(tcp_header_t))) {
		serial0_tx_PGM(PSTR("tcp: failed to prepend our header\r\n"));
		return E_NOTHING;
	}

	return E_REPLY;
}

/* generate a reset reply */
void tcp_makerst(tcp_header_t *pkt, tcp_header_t *new, uint8_t *src) {
	/* copy some header stuff from the original packet */
	new->sport = pkt->dport;
	new->dport = pkt->sport;
	/* zero-populate the SEQ and window, given we don't care */
	new->seq = 0;
	new->window = 0;
	/* increment the SEQ provided on the input packet, ACK it */
	new->ack = swap32((swap32(pkt->seq)+1));
	/* set size of packet to have no headers, and appropriate bits set */
	new->control = swap16(((sizeof(tcp_header_t)/4 << 12) | TCP_RST | TCP_ACK));
	new->crc = swap16(tcp_psuedocrc(sizeof(tcp_header_t),ip_addr,src,(uint8_t *)new));

	return;
}

/* initalise the TCP stack */
void tcp_init(void) {
	memset(&tcp_conntable,0,sizeof(tcp_conntable_t)*TCP_CONNTABLE_SIZE);
}

/* create a new connnection from the header given and source IP address */
uint8_t tcp_conntable_new(tcp_header_t pkt, uint8_t *src) {
	uint8_t n;

	for (n = 0; n < TCP_CONNTABLE_SIZE; n++) {
		if (tcp_conntable[n].state == TCP_STATE_CLOSED) {
			/* use this and return it */
			tcp_conntable[n].activity = (uint32_t) global_counters.uptime;
			tcp_conntable[n].dport = swap16(pkt.dport);
			tcp_conntable[n].sport = swap16(pkt.sport);
			memcpy(&(tcp_conntable[n].src),src,4);
			memcpy(&(tcp_conntable[n].last_seq),pkt.seq,4);
			memset(&(tcp_conntable[n].last_ack),0,4);
			tcp_conntable[n].state = TCP_STATE_SYN_REC;
			return n;
		}
	}

	return 0xff; /* failed to find a slot */
}

/* the main TCP ticker entry point */
void tcp_tick(void) {
	tcp_conntable_tick();
}

/* handle ticks on the connection table */
void tcp_conntable_tick(void) {
	uint8_t n;
	/* walk the table, if we find anything older than 10 minutes, purge it */
	for (n = 0; n < TCP_CONNTABLE_SIZE; n++) {
		if ((tcp_conntable[n].state == TCP_STATE_SYN_REC &&
			 (uint32_t) global_counters.uptime - tcp_conntable[n].activity > 10) ||
			(tcp_conntable[n].state != TCP_STATE_CLOSED &&
			 (uint32_t) global_counters.uptime - tcp_conntable[n].activity > 600)) {
			/* purge it */
			tcp_conntable[n].state = TCP_STATE_CLOSED;
#ifdef DEBUG_TCP
			serial0_tx_PGM(PSTR("tcp: purge table#"));
			serial0_tx_dec(n);
			serial0_tx_cr();
#endif
		}
	}
}

/* Compute crc of psuedoheader */
uint16_t tcp_psuedocrc(uint16_t length, uint8_t *src, uint8_t *dst, uint8_t *pkt) {
	uint32_t sum = 0;
	uint8_t n = sizeof(tcp_header_t);

	/* all this casting is needed to avoid problems with avr-gcc generating 32-bit math */

	/* psuedoheader */
	sum += (uint32_t) ((uint32_t) src[0] << 8 | (uint32_t) src[1]);
	sum += (uint32_t) ((uint32_t) src[2] << 8 | (uint32_t) src[3]);

	sum += (uint32_t) ((uint32_t) dst[0] << 8 | (uint32_t) dst[1]);
	sum += (uint32_t) ((uint32_t) dst[2] << 8 | (uint32_t) dst[3]);

	sum += ((uint32_t) 6); /* TCP protocol */
	sum += ((uint32_t) length);

	/* header */
	while (n > 1) {
		/* all this casting is needed to force avr-gcc to actually use 32 bits */
		sum += (uint32_t)(((uint32_t) * pkt << 8) | ((uint32_t) * (pkt + 1)));
		pkt += 2;
		n -= 2;
	}

	/* and the final carries */
	if (sum >> 16) {
		sum = (sum & 0xffff) + (sum >> 16);
	}

	sum = ~sum;

	return ((uint16_t)(sum));

}
